package com.ruoyi.web.controller;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;
import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.core.page.TableDataInfo;
import com.ruoyi.common.utils.StringUtils;
import com.ruoyi.system.domain.PubAichat;
import com.ruoyi.system.domain.PubAichatLog;
import com.ruoyi.system.service.IPubAichatLogService;
import com.ruoyi.system.service.IPubAichatService;
import com.ruoyi.web.utils.AIChatUtil;
import io.reactivex.Flowable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicReference;

@RestController
@RequestMapping("/web/AIChat")
public class WebAIChatController extends BaseController {

    @Autowired
    private IPubAichatService pubAichatService;
    @Autowired
    private IPubAichatLogService pubAichatLogService;
    private static final Logger logger = LoggerFactory.getLogger(WebAIChatController.class);

    // 发送信息
    @PostMapping("/sendMessage")
    public AjaxResult sendMessage(@RequestBody Map<String,Object> message ) throws NoApiKeyException, ApiException, InputRequiredException{
        String content = message.get("content").toString();
        return  AjaxResult.success().put("data", AIChatUtil.callWithMessage(content));
    }

    @GetMapping("/list/{chatId}")
    public TableDataInfo list(PubAichatLog pubAichatLog)
    {
        startPage();
        List<PubAichatLog> list = pubAichatLogService.selectPubAichatLogList(pubAichatLog);
        return getDataTable(list);
    }


    // 发送信息
    @PostMapping("/save")
    public AjaxResult save(@RequestBody PubAichat pubAichat) throws NoApiKeyException, ApiException, InputRequiredException{
        pubAichatService.insertPubAichat(pubAichat);
        return  AjaxResult.success().put("chatId", pubAichat.getId());
    }



    @GetMapping("/stream")
    public SseEmitter stream(Long userid,String content,Long chatId) {
        // 用于创建一个 SSE 连接对象
        SseEmitter emitter = new SseEmitter();
        try {
            if (StringUtils.isNotNull(content)){
                Generation gen = new Generation();
                Message userMsg = Message.builder().role(Role.USER.getValue()).content(content).build();
                new Thread(() -> {
                    try {
                        String aiResult = streamCallWithMessage(gen, userMsg, emitter);
                        streamCallWithCallback(gen, userMsg,emitter);
                        saveLog(content,aiResult,chatId);
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }).start();
            }
        } catch (Exception e) {
            // 发生错误时，关闭连接并报错
            emitter.completeWithError(e);
        }finally {

        }
        return emitter;
    }

    public void saveLog(String content, String aiResult,long chatId){
        PubAichatLog pubAichatLog = new PubAichatLog();
        pubAichatLog.setChatId(chatId);
        pubAichatLog.setContent(content);
        pubAichatLog.setIsAssistant(0l);
        pubAichatLogService.insertPubAichatLog(pubAichatLog);

        PubAichatLog pubAichatLogs = new PubAichatLog();
        pubAichatLogs.setContent(aiResult);
        pubAichatLogs.setIsAssistant(1l);
        pubAichatLogService.insertPubAichatLog(pubAichatLogs);
    }


    public  String streamCallWithMessage(Generation gen, Message userMsg,SseEmitter emitter)
            throws NoApiKeyException, ApiException, InputRequiredException,IOException {
        GenerationParam param = buildGenerationParam(userMsg);
        Flowable<GenerationResult> result = gen.streamCall(param);
        StringBuilder fullContent = new StringBuilder();
        result.blockingForEach(message -> handleGenerationResult(message, fullContent,emitter));
        return fullContent.toString();
    }

    /**
     * 流式输出
     * @param message
     * @param fullContent
     */
    private  void handleGenerationResult(GenerationResult message, StringBuilder fullContent,SseEmitter emitter) throws IOException {
        fullContent.append(message.getOutput().getChoices().get(0).getMessage().getContent());
        // 发送消息
        emitter.send(SseEmitter.event().name("message").data(JsonUtils.toJson(message)));
        String finishReason = message.getOutput().getChoices().get(0).getFinishReason();
        if ("stop".equals(finishReason)) {
            // 完成发送
            emitter.complete();
        }
    }

    public  void streamCallWithCallback(Generation gen, Message userMsg,SseEmitter emitter)
            throws NoApiKeyException, ApiException, InputRequiredException, InterruptedException {
        GenerationParam param = buildGenerationParam(userMsg);
        Semaphore semaphore = new Semaphore(0);
        StringBuilder fullContent = new StringBuilder();

        gen.streamCall(param, new ResultCallback<GenerationResult>() {
            @Override
            public void onEvent(GenerationResult message) {
                try {
                    handleGenerationResult(message, fullContent,emitter);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

            @Override
            public void onError(Exception err) {
                semaphore.release();
            }

            @Override
            public void onComplete() {
                semaphore.release();
            }
        });

        semaphore.acquire();
    }
    private  GenerationParam buildGenerationParam(Message userMsg) {
        return GenerationParam.builder()
                .model("qwen-turbo")
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }

}